import streamlit as st
from PyPDF2 import PdfReader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_core.prompts import ChatPromptTemplate
from langchain_community.vectorstores import FAISS
from langchain.tools.retriever import create_retriever_tool
from langchain.agents import AgentExecutor, create_tool_calling_agent
from langchain_community.embeddings import DashScopeEmbeddings
from langchain.chat_models import init_chat_model
import os
from dotenv import load_dotenv
load_dotenv(override=True)

DeepSeek_API_KEY = os.getenv("DEEPSEEK_API_KEY")
dashscope_api_key = os.getenv("DASHSCOPE_API_KEY")
os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"
embeddings = DashScopeEmbeddings(model="text-embedding-v1", dashscope_api_key="")

def pdf_read(pdf_doc):
    text = ""
    for pdf in pdf_doc:
        pdf_reader = PdfReader(pdf)
        for page in pdf_reader.pages:
            text += page.extract_text()
    return text

def get_chunks(text):
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
    chunks = text_splitter.split_text(text)
    return chunks

def vector_store(text_chunks):
    vector_store = FAISS.from_texts(text_chunks, embedding=embeddings)
    vector_store.save_local("faiss_db")

def get_conversational_chain(tools, ques):
    llm = init_chat_model(model="deepseek-chat", model_provider="deepseek")
    prompt = ChatPromptTemplate.from_messages(
        [("system", """你是一个AI助手，请根据提供的上下文信息回答用户的问题，回答必须基于上述上下文，不要编造信息。如果上下文中没有足够的信息，请明确回答：'根据提供的信息无法回答该问题。'"""),
         ("placeholder", "{chat_history}"),
         ("human", "{input}"),
         ("placeholder", "{agent_scratchpad}"),
    ])
    tool = [tools]
    agent = create_tool_calling_agent(llm, tool, prompt)
    agent_executor = AgentExecutor(agent=agent, tools=tool, verbose=True)
    response = agent_executor.invoke({"input": ques})
    print(response)
    st.write("回答:", response['output'])

def check_database_exists():
    return os.path.exists("faiss_db") and os.path.exists("faiss_db/index.faiss")

def user_input(user_question):
    if not check_database_exists():
        st.error("请先上传PDF文件并点击")
        st.info("步骤1:上传PDF, 步骤2:点击处理, 步骤3:开始提问")
        return
    try:
        new_db = FAISS.load_local("faiss_db", embeddings=embeddings, allow_dangerous_deserialization=True)
        retriever = new_db.as_retriever()
        retrieval_chain = create_retriever_tool(retriever, "pdf_extractor", "")
        get_conversational_chain(retrieval_chain, user_question)

    except Exception as e:
        st.error(f"加载数据库时出错{str(e)}")
        st.info("请重新处理PDF文件")

def main():
    st.set_page_config("langchain自学")
    st.header("langchain自学")
    col1, col2 = st.columns([3, 1])
    with col1:
        if check_database_exists():
            pass
        else:
            st.warning("请先上传并处理PDF文件")
    with col2:
        if st.button("清除数据库"):
            try:
                import shutil
                if os.path.exists("faiss_db"):
                    shutil.rmtree("faiss_db")
                st.success("数据库已经清楚")
                st.rerun()
            except Exception as e:
                st.error(f"清除失败: {e}")

    user_question = st.text_input("请输入问题", placeholder="例如,这个文档的主要内容是什么?", disabled=not check_database_exists())
    if user_question:
        if check_database_exists():
            with st.spinner("AI正在分析文档..."):
                user_input(user_question)
        else:
            st.error("请先上传并处理PDF文件")
    with st.sidebar:
        st.title("文档管理")
        if check_database_exists():
            st.success("数据库状态: 已就绪")
        else:
            st.info("状态: 等待上传PDF")
        st.markdown("---")
        pdf_doc = st.file_uploader(
            "上传PDF",
            accept_multiple_files=True,
            type=['pdf'],
            help="支持上传多PDF文件"
        )
        if pdf_doc:
            st.info(f"已选择{len(pdf_doc)}个文件")
            for i, pdf in enumerate(pdf_doc, 1):
                st.write(f"{i}. {pdf.name}")

        process_button = st.button(
            "处理文件",
            disabled=not pdf_doc,
            use_container_width=True
        )
        if process_button:
            if pdf_doc:
                with st.spinner("正在处理PDF文件"):
                    try:
                        raw_text = pdf_read(pdf_doc)
                        if not raw_text.strip():
                            st.error("无法从pdf中提取文本,请检测文件是否有效")
                            return
                        text_chunks = get_chunks(raw_text)
                        st.info(f"文本已分割为{len(text_chunks)}个文件")

                        vector_store(text_chunks)
                        st.success("PDF处理完成,可以开始提问")
                        st.balloons()
                        st.rerun()
                    except Exception as e:
                        st.error(f"处理PDF时出错: {str(e)}")
            else:
                st.warning("请先选择PDF文件")

        with st.expander("使用说明"):
            st.markdown("""
            ### 📝 使用步骤
            1. **上传文件**  
               点击上传区域，选择一个或多个 PDF 文件。
            2. **处理文件**  
               点击“处理文件”按钮，系统将提取内容并构建知识库。
            3. **提问**  
               转到主页，在输入框中提出您的问题。
            4. **获取答案**  
               AI 将基于您上传的 PDF 内容生成回答。

            ### 💡 提示
            - ✅ 支持同时上传多个 PDF 文件  
            - ⏳ 处理大文件或较多文件时可能需要等待片刻  
            - 🔄 可随时点击“清除数据库”重新开始
            """)

if __name__ == "__main__":
    main()